from osgeo import gdal
import os
import numpy as np
import warnings

warnings.filterwarnings(action='ignore')

land_list = [10, 11, 12, 13, 19, 94]
irrigated_list = [1, 2, 3]
rainfed_list = [4, 5, 6, 7, 8]


def kth_smallest(matrix, k):
    tmp = []
    for i in matrix:
        tmp.extend(i)
    sort_data = np.sort(list(set(tmp)))
    if len(sort_data) == 1:
        return sort_data[0]
    else:
        return sort_data[k - 1]


def min_max_normalization_fill(im_data, k, fill):
    second_min = kth_smallest(im_data, k)
    max = np.max(im_data)
    if second_min == max:
        out_data = im_data
    else:
        max_min_data = (im_data - second_min) / (max - second_min)
        out_data = max_min_data * 2.0 - 1.0
    return np.where(out_data < -1, fill, out_data)


def raster_copy_with_nodata(s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
                            t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
                            nodata):
    """Copy a band of raster into the output file with nodata values.

       Function copied from gdal_merge.py
    """
    try:
        import numpy as Numeric
    except ImportError:
        import Numeric

    s_band = s_fh.GetRasterBand(s_band_n)
    t_band = t_fh.GetRasterBand(t_band_n)
    data_src = s_band.ReadAsArray(s_xoff, s_yoff, s_xsize, s_ysize,
                                  t_xsize, t_ysize)
    # data_src = min_max_normalization_fill(data_src, 2, nodata)
    data_dst = t_band.ReadAsArray(t_xoff, t_yoff, t_xsize, t_ysize)
    nodata_test = Numeric.equal(data_src, nodata)
    to_write = Numeric.choose(nodata_test, (data_src, data_dst))
    t_band.WriteArray(to_write, t_xoff, t_yoff)

    return 0


def raster_copy(s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
                t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
                nodata=None):
    """Copy a band of raster into the output file.

       Function copied from gdal_merge.py
    """
    if nodata is not None:
        return raster_copy_with_nodata(s_fh, s_xoff, s_yoff, s_xsize, s_ysize,
                                       s_band_n, t_fh, t_xoff, t_yoff, t_xsize,
                                       t_ysize, t_band_n, nodata)
    s_band = s_fh.GetRasterBand(s_band_n)
    t_band = t_fh.GetRasterBand(t_band_n)
    data = s_band.ReadRaster(s_xoff, s_yoff, s_xsize, s_ysize,
                             t_xsize, t_ysize, t_band.DataType)
    t_band.WriteRaster(t_xoff, t_yoff, t_xsize, t_ysize, data, t_xsize,
                       t_ysize, t_band.DataType)
    return 0


def copy_into(t_fh, in_ds, s_band=1, t_band=1, nodata_arg=None):
    """Copy this files image into target file.
    This method will compute the overlap area of the file_info objects
    file, and the target gdal.Dataset object, and copy the image data
    for the common window area.  It is assumed that the files are in
    a compatible projection. no checking or warping is done.  However,
    if the destination file is a different resolution, or different
    image pixel type, the appropriate resampling and conversions will
    be done (using normal GDAL promotion/demotion rules).

    :param t_fh: gdal.Dataset object for the file into which some or all
                 of this file may be copied.
    :param s_band:
    :param t_band:
    :param nodata_arg:

    :return: 1 on success (or if nothing needs to be copied), and zero one
             failure.

    """
    geotransform = in_ds.GetGeoTransform()
    xsize = in_ds.RasterXSize
    ysize = in_ds.RasterYSize
    ulx = geotransform[0]
    uly = geotransform[3]
    lrx = ulx + geotransform[1] * xsize
    lry = uly + geotransform[5] * ysize
    t_geotransform = t_fh.GetGeoTransform()
    t_ulx = t_geotransform[0]
    t_uly = t_geotransform[3]
    t_lrx = t_geotransform[0] + t_fh.RasterXSize * t_geotransform[1]
    t_lry = t_geotransform[3] + t_fh.RasterYSize * t_geotransform[5]

    # figure out intersection region
    tgw_ulx = max(t_ulx, ulx)
    tgw_lrx = min(t_lrx, lrx)
    if t_geotransform[5] < 0:
        tgw_uly = min(t_uly, uly)
        tgw_lry = max(t_lry, lry)
    else:
        tgw_uly = max(t_uly, uly)
        tgw_lry = min(t_lry, lry)

    # do they even intersect?
    if tgw_ulx >= tgw_lrx:
        return 1
    if t_geotransform[5] < 0 and tgw_uly <= tgw_lry:
        return 1
    if t_geotransform[5] > 0 and tgw_uly >= tgw_lry:
        return 1

    # compute target window in pixel coordinates.
    tw_xoff = int((tgw_ulx - t_geotransform[0]) / t_geotransform[1] + 0.1)
    tw_yoff = int((tgw_uly - t_geotransform[3]) / t_geotransform[5] + 0.1)
    tw_xsize = int((tgw_lrx - t_geotransform[0]) / t_geotransform[1] + 0.5) - tw_xoff
    tw_ysize = int((tgw_lry - t_geotransform[3]) / t_geotransform[5] + 0.5) - tw_yoff

    if tw_xsize < 1 or tw_ysize < 1:
        return 1

    # Compute source window in pixel coordinates.
    sw_xoff = int((tgw_ulx - geotransform[0]) / geotransform[1])
    sw_yoff = int((tgw_uly - geotransform[3]) / geotransform[5])
    sw_xsize = int((tgw_lrx - geotransform[0])
                   / geotransform[1] + 0.5) - sw_xoff
    sw_ysize = int((tgw_lry - geotransform[3])
                   / geotransform[5] + 0.5) - sw_yoff

    if sw_xsize < 1 or sw_ysize < 1:
        return 1
    return \
        raster_copy(in_ds, sw_xoff, sw_yoff, sw_xsize, sw_ysize, s_band,
                    t_fh, tw_xoff, tw_yoff, tw_xsize, tw_ysize, t_band,
                    nodata_arg)


def get_extent(in_fn):
    ds = gdal.Open(in_fn)
    geotrans = list(ds.GetGeoTransform())
    xsize = ds.RasterXSize
    ysize = ds.RasterYSize
    min_x = geotrans[0]
    max_y = geotrans[3]
    max_x = geotrans[0] + xsize * geotrans[1]
    min_y = geotrans[3] + ysize * geotrans[5]
    ds = None
    return min_x, max_y, max_x, min_y


def find_layers_files(dir_name, format):
    file_list = []
    for home, dirs, files in os.walk(dir_name):
        for filename in sorted(files):
            if filename.endswith(format):
                file_list.append(os.path.join(home, filename))
    return file_list


def del_file(path):
    for i in os.listdir(path):
        path_file = os.path.join(path, i)
        if os.path.isfile(path_file):
            os.remove(path_file)
        else:
            del_file(path_file)


def set_color_table():
    ct = gdal.ColorTable()
    ctmap = {0: (0, 0, 0, 255), 1: (122, 114, 238, 255), 2: (30, 38, 130, 255), 3: (166, 252, 168, 255),
             4: (0, 134, 0, 255), 5: (16, 146, 26, 255), 6: (252, 244, 100, 255), 7: (200, 200, 2, 255),
             8: (140, 140, 0, 255), 9: (192, 192, 192, 255), 10: (163, 255, 115, 255), 11: (238, 2, 48, 255),
             12: (212, 142, 254, 255), 13: (170, 36, 250, 255), 14: (0, 234, 0, 255), 15: (122, 114, 238, 255)}
    for i in range(256):
        index = i
        index = min(index, len(ctmap) - 1)
        ct.SetColorEntry(i, ctmap[index])
    return ct
